'''
增强版兰顿蚂蚁教程：有三只不同的蚂蚁。当速度不是水平竖直的时候，兰顿蚂蚁的运动状况如何？
作者：侯展意
协议：木兰2.0
'''

import numpy as np
import takagiabm as tak
import random


def agentSetup(model):
    colorList = ['#ff0000', '#00ff00', '#0000ff']
    for i in range(3):
        a = tak.GridAgent(np.array([int(model.width/2),int(model.height/2)]))
        triangle2 = tak.scale(tak.takTriangle, 3)  # 设置形状为takagi库的三角形，缩放比例为3.
        a.setShape(triangle2)  # 设置形状为三角形。
        a.setSpeed(np.array([1,0]))  # 设置速度。只有第一只蚂蚁速度是水平数值的，因此会有不同的运动规律。
        a.setColor(colorList[i%3])  # 设置颜色
        model.addAgent(a)


def agentLoop(agent):  # 代理人(也就是蚂蚁)执行的函数
    global stat
    color = agent.getCellColor()

    if (color == '#ffffff'):
        agent.setCellColor('#000000')
        agent.turn('RIGHT')
    else:
        agent.setCellColor('#ffffff')
        agent.turn('LEFT')
    agent.move()


stat = 0
tak.GridAgent = tak.prepareAgent(tak.GridAgent,agentLoop)
tak.GridModel= tak.prepareModel(tak.GridModel,modelInitFunc=agentSetup,width=70,height=70,
                                agentActivationPolicy='casual')
if __name__ == "__main__":
    tak.simStart(__file__, tak.GridModel)  # 这个方法同时启动了图形界面。
